/*
 * tcp_upgrade.c
 *
 *  Created on: Mar 21, 2021
 *      Author: kychu
 */

#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <stdint.h>
#include <dirent.h>
#include <sys/param.h>
#include <sys/stat.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_err.h"
#include "esp_vfs.h"
#include "nvs_flash.h"
#include "esp_netif.h"
#include "esp_idf_version.h"

#include "lwip/err.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include <lwip/netdb.h>

#include "esp_ota_ops.h"
#include "esp_flash_partitions.h"
#include "esp_partition.h"
#include "nvs.h"
#include "errno.h"

static const char *TAG = "OTA";

static const char *cmd_ok = "OK!";
static const char *cmd_err = "ERROR!";

int up_notification(const int sock, const char *cmd, const char *desc);
int tcp_upgrade(const int sock, size_t size, uint8_t *buff, size_t cache_len)
{
  esp_err_t err;
  int ret = 0;
  /* update handle : set by esp_ota_begin(), must be freed via esp_ota_end() */
  esp_ota_handle_t update_handle = 0 ;
  const esp_partition_t *update_partition = NULL;

  if(size == 0) {
    ESP_LOGE(TAG, "binary size error");
    up_notification(sock, cmd_err, "size error");
    return -1;
  }
  const esp_partition_t *configured = esp_ota_get_boot_partition();
  const esp_partition_t *running = esp_ota_get_running_partition();

  if (configured != running) {
    ESP_LOGW(TAG, "Configured OTA boot partition at offset 0x%08x, but running from offset 0x%08x",
                 configured->address, running->address);
    ESP_LOGW(TAG, "(This can happen if either the OTA boot data or preferred boot image become corrupted somehow.)");
  }
  ESP_LOGI(TAG, "Running partition type %d subtype %d (offset 0x%08x)",
             running->type, running->subtype, running->address);

  update_partition = esp_ota_get_next_update_partition(NULL);
  ESP_LOGI(TAG, "Writing to partition subtype %d at offset 0x%x",
               update_partition->subtype, update_partition->address);
//  assert(update_partition != NULL);
  if(update_partition == NULL) {
    ESP_LOGE(TAG, "get update partition failed");
    up_notification(sock, cmd_err, "partition error");
    return -1;
  }
  if(up_notification(sock, cmd_ok, "transefer start") < 0) { // start transfer
    return -1;
  }
  int binary_file_length = 0;
  /*deal with all receive packet*/
  bool image_header_was_checked = false;
  do {
    int data_read = recv(sock, buff, cache_len, 0);
    if (data_read < 0) {
      ESP_LOGE(TAG, "Error occurred during receiving: errno %d", errno);
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 4, 0)
      if(image_header_was_checked) {
        esp_ota_abort(update_handle);
      }
#else
      esp_ota_end(update_handle);
#endif
      up_notification(sock, cmd_err, "RECV failed");
      ret = -1; goto exit;
    } else if(data_read == 0) {
      if (errno == ECONNRESET || errno == ENOTCONN) {
        ESP_LOGE(TAG, "Connection closed, errno = %d", errno);
      }
      break;
    } else {
      if (image_header_was_checked == false) {
        esp_app_desc_t new_app_info;
        if (data_read > sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t) + sizeof(esp_app_desc_t)) {
          // check current version with downloading
          memcpy(&new_app_info, &buff[sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t)], sizeof(esp_app_desc_t));
          ESP_LOGI(TAG, "New firmware version: %s", new_app_info.version);

          esp_app_desc_t running_app_info;
          if (esp_ota_get_partition_description(running, &running_app_info) == ESP_OK) {
            ESP_LOGI(TAG, "Running firmware version: %s", running_app_info.version);
          }

          const esp_partition_t* last_invalid_app = esp_ota_get_last_invalid_partition();
          esp_app_desc_t invalid_app_info;
          if (esp_ota_get_partition_description(last_invalid_app, &invalid_app_info) == ESP_OK) {
            ESP_LOGI(TAG, "Last invalid firmware version: %s", invalid_app_info.version);
          }

          // check current version with last invalid partition
          if (last_invalid_app != NULL) {
            if (memcmp(invalid_app_info.version, new_app_info.version, sizeof(new_app_info.version)) == 0) {
              ESP_LOGW(TAG, "New version is the same as invalid version.");
              ESP_LOGW(TAG, "Previously, there was an attempt to launch the firmware with %s version, but it failed.", invalid_app_info.version);
              ESP_LOGW(TAG, "The firmware has been rolled back to the previous version.");
              up_notification(sock, cmd_err, "same as INVALID version");
              ret = -1; goto exit;
            }
          }

          if (memcmp(new_app_info.version, running_app_info.version, sizeof(new_app_info.version)) == 0) {
            ESP_LOGW(TAG, "Current running version is the same as a new. We will not continue the update.");
            if(send(sock, cmd_err, strlen(cmd_err), 0) < 0) { // try to notify host
              ESP_LOGE(TAG, "Error occurred during sending: errno %d", errno);
            }
            up_notification(sock, cmd_err, "same as CURRENT version");
            ret = -1; goto exit;
          }

          image_header_was_checked = true;
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 4, 0)
          err = esp_ota_begin(update_partition, OTA_WITH_SEQUENTIAL_WRITES, &update_handle);
#else
          err = esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);
#endif
          if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 4, 0)
            esp_ota_abort(update_handle);
#else
            esp_ota_end(update_handle);
#endif
            up_notification(sock, cmd_err, "OTA begin failed");
            ret = -1; goto exit;
          }
          ESP_LOGI(TAG, "esp_ota_begin succeeded");
        } else {
          ESP_LOGE(TAG, "received package is not fit len");
          up_notification(sock, cmd_err, "package too short");
          ret = -1; goto exit;
        }
      }
      err = esp_ota_write( update_handle, (const void *)buff, data_read);
      if(err != ESP_OK) {
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 4, 0)
        esp_ota_abort(update_handle);
#else
        esp_ota_end(update_handle);
#endif
        up_notification(sock, cmd_err, "OTA write failed");
        ret = -1; goto exit;
      }
      binary_file_length += data_read;
      ESP_LOGD(TAG, "Written image length %d", binary_file_length);
    }
  } while(binary_file_length < size);
  ESP_LOGI(TAG, "Total Write: %d, Binary size: %d", binary_file_length, size);
  if(binary_file_length != size) {
    ESP_LOGE(TAG, "Error in receiving complete file");
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 4, 0)
    esp_ota_abort(update_handle);
#else
    esp_ota_end(update_handle);
#endif
    up_notification(sock, cmd_err, "Connection closed");
    ret = -1; goto exit;
  }
  err = esp_ota_end(update_handle);
  if (err != ESP_OK) {
    if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
      ESP_LOGE(TAG, "Image validation failed, image is corrupted");
      up_notification(sock, cmd_err, "Image validation failed");
    } else {
      up_notification(sock, cmd_err, "OTA end failed");
    }
    ESP_LOGE(TAG, "esp_ota_end failed (%s)!", esp_err_to_name(err));
    ret = -1; goto exit;
  }

  err = esp_ota_set_boot_partition(update_partition);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_set_boot_partition failed (%s)!", esp_err_to_name(err));
    up_notification(sock, cmd_err, "set boot partition failed");
    ret = -1; goto exit;
  }
  up_notification(sock, cmd_ok, "OTA succeed");
  vTaskDelay(pdMS_TO_TICKS(500));
  ESP_LOGI(TAG, "Prepare to restart system!");
  esp_restart();

exit:
  return ret;
}
